package it.unibo.refolding.alg;

import it.unibo.refolding.common.Callback;

import java.util.Iterator;

import org.apache.commons.lang3.builder.ReflectionToStringBuilder;
import org.apache.commons.lang3.builder.ToStringStyle;
import org.apache.commons.lang3.mutable.MutableInt;

import com.google.common.base.Preconditions;
import com.google.common.collect.AbstractSequentialIterator;

/**
 * @see http://lwn.net/Articles/184495/
 * Based on the augmented version of the red-black tree implementation of the linux kernel found in libosl.
 */
public class RbTree<T extends Comparable<T>> {

	private final RbRoot<T> root = new RbRoot<T>();
    private final Callback<RbNode<T>> sizeCallback = new Callback<RbNode<T>>() {
        @Override
        public void call(RbNode<T> t) {
            t.size = t.leftSize() + 1 + t.rightSize();
        }
    };
	
	private RbNode<T> insertAttempt(RbNode<T> node) {
	    T key = node.key;
		RbNode<T> p = root.node;
		RbNode<T> parent = null;
		int cmp = 0;
		while (p != null) {
			parent = p;
			cmp = key.compareTo(parent.key);
            if (cmp < 0) {
				p = p.left;
			} else if (cmp > 0) {
				p = p.right;
			} else {
				return parent;
			}
		}
		node.linkNode(parent);
		// Next code block is also part of the linking procedure.
		if (parent == null) {
		    root.node = node;
		} else {
		    if (cmp < 0) {
		        assert parent.left == null;
		        parent.left = node;
		    } else if (cmp > 0) {
		        assert parent.right == null;
		        parent.right = node;
		    } else {
		        throw new IllegalStateException("impossible");
		    }
		}
		return node;
	}
	
	private RbNode<T> findOrInsert(RbNode<T> node) {
        Preconditions.checkArgument(node.parent == null & !node.black & node.left == null & node.right == null & node.size == 0);
		RbNode<T> attempt = insertAttempt(node);
		if (attempt != node) {
		    // Already found.
			return attempt;
		} else {
		    // Inserted.
			RbNode.insertColor(attempt, root);
	        RbNode.augmentInsert(node, sizeCallback);
			return node;
		}
	}

	public RbNode<T> insert(T key) {
	    RbNode<T> node = new RbNode<>(key);
	    RbNode<T> actual = findOrInsert(node);
	    if (actual == node) {
	        return node;
	    } else {
	        return null;
	    }
	}
	
	public RbNode<T> search(T key) {
	    RbNode<T> n = root.node;
        while (n != null) {
            if (key.compareTo(n.key) < 0) {
                n = n.left;
            } else if (key.compareTo(n.key) > 0) {
                n = n.right;
            } else {
                return n;
            }
        }
        return null;
	}
	
	public int rank(RbNode<T> node) {
	    MutableInt rank = new MutableInt();
	    RbNode<T> rootNode = RbNode.rank(node, rank);
	    Preconditions.checkArgument(
	        rootNode == root.node, "Node %s with root %s is not from this tree with root %s.",
	        node, rootNode, root.node);
	    return rank.intValue();
	}
	
	public int rank(T key) {
	    // For now the key must exist in the tree.
	    RbNode<T> node = search(key);
	    if (node == null) {
	        return 0;
	    }
	    return rank(node);
	}
	
	public RbNode<T> delete(T key) {
        RbNode<T> node = search(key);
        if (node == null) {
            return null;
        }
	    delete(node);
	    return node;
	}

    public void delete(RbNode<T> node) {
        // See http://lxr.free-electrons.com/source/arch/x86/mm/pat_rbtree.c.
        // We should check first if the node is in our tree...
        RbNode<T> deepest = RbNode.augmentEraseBegin(node);
        RbNode.erase(node, root);
        RbNode.augmentEraseEnd(deepest, sizeCallback);
    }

    public Iterable<RbNode<T>> nodeIterable() {
        return new Iterable<RbNode<T>>() {
            @Override
            public Iterator<RbNode<T>> iterator() {
                return nodeIterator();
            }
        };
    }
    
    public Iterator<RbNode<T>> nodeIterator() {
        return new AbstractSequentialIterator<RbNode<T>>(RbNode.first(root)) {
            @Override
            protected RbNode<T> computeNext(RbNode<T> previous) {
                return RbNode.next(previous);
            }
        };
    }

    public Iterable<T> keyIterable() {
        return new Iterable<T>() {
            @Override
            public Iterator<T> iterator() {
                return keyIterator();
            }
        };
    }
    
    public Iterator<T> keyIterator() {
        final Iterator<RbNode<T>> nodeIterator = nodeIterator();
        return new Iterator<T>() {
            @Override
            public boolean hasNext() {
                return nodeIterator.hasNext();
            }

            @Override
            public T next() {
                return nodeIterator.next().key;
            }

            @Override
            public void remove() {
                nodeIterator.remove();
            }
        };
    }

    public int size() {
        if (isEmpty()) {
            return 0;
        } else {
            RbNode<T> last = RbNode.last(root);
            int rank = rank(last);
            assert rank >= 1;
            return rank;
        }
    }

    @Override
    public String toString() {
        return ReflectionToStringBuilder.toString(this, ToStringStyle.SHORT_PREFIX_STYLE);
    }

    public boolean isEmpty() {
        return root.node == null;
    }
}
